package ai

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"time"

	"anders-cli/internal/config"

	"github.com/go-resty/resty/v2"
)

// XAIBackend xAI AI后端实现
type XAIBackend struct {
	client  *resty.Client
	baseURL string
	model   string
	apiKey  string
}

// NewXAIBackend 创建新的xAI后端实例
func NewXAIBackend() *XAIBackend {
	cfg := config.Get()

	return &XAIBackend{
		client:  resty.New(),
		baseURL: cfg.XAI.BaseURL,
		model:   cfg.XAI.Model,
		apiKey:  os.Getenv(cfg.XAI.ApiKeyEnv),
	}
}

// Name 返回后端名称
func (b *XAIBackend) Name() string {
	return "xAI"
}

// GetModel 返回当前使用的模型名称
func (b *XAIBackend) GetModel() string {
	return b.model
}

// Query 执行AI查询
func (b *XAIBackend) Query(prompt string, tools []Tool) (string, error) {
	// 从配置获取超时时间
	cfg := config.Get()
	queryTimeout := time.Duration(cfg.MultiAgent.QueryTimeout) * time.Second
	if queryTimeout <= 0 {
		queryTimeout = 120 * time.Second // 默认120秒
	}

	// 创建超时上下文
	ctx, cancel := context.WithTimeout(context.Background(), queryTimeout)
	defer cancel()

	// 查询开始

	// 检查API密钥
	if b.apiKey == "" {
		return "", fmt.Errorf("xAI API密钥未设置，请设置环境变量 %s", config.Get().XAI.ApiKeyEnv)
	}

	// 构建请求消息
	messages := []Message{
		{
			Role:    "user",
			Content: prompt,
		},
	}

	// 构建请求体
	requestBody := map[string]interface{}{
		"model":    b.model,
		"messages": messages,
	}

	// 如果有工具，添加到请求中
	if len(tools) > 0 {
		requestBody["tools"] = tools
		requestBody["tool_choice"] = "auto"
	}

	// 发送请求（带超时上下文）
	resp, err := b.client.R().
		SetHeader("Content-Type", "application/json").
		SetHeader("Authorization", "Bearer "+b.apiKey).
		SetBody(requestBody).
		SetContext(ctx).
		Post(b.baseURL + "/chat/completions")

	if err != nil {
		return "", fmt.Errorf("xAI 请求失败: %v", err)
	}

	if resp.StatusCode() != 200 {
		return "", fmt.Errorf("xAI API 返回错误状态码: %d", resp.StatusCode())
	}

	// 解析响应
	var response struct {
		Choices []struct {
			Message struct {
				Content   string `json:"content"`
				ToolCalls []struct {
					Function struct {
						Name      string          `json:"name"`
						Arguments json.RawMessage `json:"arguments"`
					} `json:"function"`
				} `json:"tool_calls,omitempty"`
			} `json:"message"`
		} `json:"choices"`
	}

	if err := json.Unmarshal(resp.Body(), &response); err != nil {
		return "", fmt.Errorf("解析响应失败: %v", err)
	}

	if len(response.Choices) == 0 {
		return "", fmt.Errorf("未收到有效响应")
	}

	content := response.Choices[0].Message.Content

	return content, nil
}
